#!/usr/bin/python
import serial
import time
import struct
import math

import rospy
import tf
from geometry_msgs.msg import Twist, Quaternion
from nav_msgs.msg import Odometry

class SweepbotDriver(): 
    def __init__(self, node_name):
        rospy.init_node(node_name,disable_signals=True)
        self.serialport_name = rospy.get_param('~serialport_name', default='/dev/sweepbot')
        self.baudrate = rospy.get_param('~baudrate', default=115200)
        self.linear_coef = rospy.get_param('~linear_coef', default=128)
        self.angular_coef = rospy.get_param('~angular_coef', default=42)
        self.left_coef = rospy.get_param('~left_coef', default=1)
        self.right_coef = rospy.get_param('~right_coef', default =1)
        self.encoder_ticks_per_rev = rospy.get_param('~encoder_ticks_per_rev', default=600)
        self.base_width = rospy.get_param('~base_width', default=0.33)
        self.wheel_diameter = rospy.get_param('~wheel_diameter', default=0.125)
        self.pub_odom = rospy.Publisher('/odom', Odometry, queue_size=10)
        self.sub_vel = rospy.Subscriber('/cmd_vel', Twist, self.vel_callback, queue_size=1)
        self.tf_broadcaster = tf.TransformBroadcaster()
        self.odom = Odometry()
        self.odom.header.frame_id = 'odom'
        self.odom.child_frame_id = 'base_link'
        self.time_prev = rospy.Time.now() - rospy.Time.from_sec(1)
        self.x = 0
        self.y = 0
        self.theta = 0
        self.encoders = [0, 0]
        self.encoders_sum = [0, 0]
        self.direction = -1
        self.conn = serial.Serial(self.serialport_name, self.baudrate, timeout=1)

    def reconnect(self, side):
        print('{} {} reconnecting...'.format(side, time.time()))
        while True:
            try:
                self.conn = serial.Serial(self.serialport_name, self.baudrate, timeout=1)
            except:
                time.sleep(1)
                print('{} {} reconnecting...'.format(side, time.time()))
            else:
                print('{} {} reconnected!'.format(side, time.time()))
                break

    def set_speed(self, v1, v2):
        v1 = int(v1)
        v2 = int(v2)
        d1 = 1
        d2 = 1
        if v1 < 0:
            d1 = 0
            v1 = -v1
        if v2 < 0:
            d2 = 0
            v2 = -v2
        v1 = v1 if v1 < 128 else 128
        v2 = v2 if v2 < 128 else 128
        msg = [0xff, 0xfe, v2, v1, d2, d1, 0x00, 0x00, 0x00, 0x00]
        #print('{:.2f} send: {}'.format(time.time(), msg))
        self.conn.write(msg)
        #while True:
        #    try:
        #        self.conn.write(msg)
        #        break
        #    except serial.serialutil.SerialException:
        #        self.reconnect('send')

    def stop(self):
        print('stop')
        self.conn.close()

    def read(self):
        #try:
        msg = self.conn.read(8)
        #except:
        #    self.reconnect('read')
        #    return
        #print('{:.2f} recv: {}'.format(time.time(), [ord(m) for m in msg]))
        if ord(msg[0]) == 0xff and ord(msg[1])== 0xfe:
            if ord(msg[5]) == 0x02:
                self.encoders[0] = -ord(msg[4])
            else:
                self.encoders[0] = ord(msg[4])
            if ord(msg[3]) == 0x02:
                self.encoders[1] = -ord(msg[2])
            else:
                self.encoders[1] = ord(msg[2])
        else:
            self.encoders = [0, 0]
            print('{} received message invalid: {}'.format(time.time(), msg))
        self.encoders_sum[0] += self.encoders[0]
        self.encoders_sum[1] += self.encoders[1]

    def vel_callback(self, msg):
        v1 = self.linear_coef * msg.linear.x
        v2 = self.linear_coef * msg.linear.x
        v1 -= self.angular_coef * msg.angular.z
        v2 += self.angular_coef * msg.angular.z
        v1 *= self.left_coef
        v2 *= self.right_coef
        #set speed
        self.set_speed(v1, v2)

    def update(self):
        #get encoders
        self.read()
        #calculate
        time_current = rospy.Time.now()
        time_elapsed = (time_current - self.time_prev).to_sec()
        if time_elapsed == 0:
            time_elapsed = 1e-3
        self.time_prev = time_current
        dleft = self.left_coef * math.pi * self.wheel_diameter * \
                self.encoders[0] / self.encoder_ticks_per_rev
        dright = self.right_coef * math.pi * self.wheel_diameter * \
                self.encoders[1] / self.encoder_ticks_per_rev
        d = (dleft + dright) / 2
        dtheta = (dright - dleft) / self.base_width
        if d != 0:
            dx = math.cos(dtheta) * d
            dy = math.sin(dtheta) * d
            self.x += dx*math.cos(self.theta)-dy*math.sin(self.theta)
            self.y += dx*math.sin(self.theta)+dy*math.cos(self.theta)
        self.theta += dtheta

        self.odom.header.stamp = time_current
        self.odom.pose.pose.position.x = self.x
        self.odom.pose.pose.position.y = self.y
        q = tf.transformations.quaternion_from_euler(0,0,self.theta)
        self.odom.pose.pose.orientation.x = q[0]
        self.odom.pose.pose.orientation.y = q[1]
        self.odom.pose.pose.orientation.z = q[2]
        self.odom.pose.pose.orientation.w = q[3]
        self.odom.twist.twist.linear.x = d / time_elapsed
        self.odom.twist.twist.angular.z = dtheta / time_elapsed

    def run(self):
        while True:
            try:
                self.update()
                self.pub_odom.publish(self.odom)
                self.tf_broadcaster.sendTransform((self.x,self.y,0),
                    tf.transformations.quaternion_from_euler(0, 0, self.theta),
                    rospy.Time.now(), 'base_link', 'odom')
            except KeyboardInterrupt:
                print('exit.')
                self.stop()
                break

if __name__ == '__main__':
    sweepbot_driver = SweepbotDriver('sweepbot_driver_node')
    sweepbot_driver.run()

